import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import SAGEConv

class CompleteNetworkGNNModel(nn.Module):
    def __init__(self, input_dim, hidden_dims=[256, 128, 64], output_dim=1, dropout=0.3):
        super(CompleteNetworkGNNModel, self).__init__()
        self.layers = nn.ModuleList()
        self.batch_norms = nn.ModuleList()
        self.layers.append(SAGEConv(input_dim, hidden_dims[0]))
        self.batch_norms.append(nn.BatchNorm1d(hidden_dims[0]))
        for i in range(1, len(hidden_dims)):
            self.layers.append(SAGEConv(hidden_dims[i - 1], hidden_dims[i]))
            self.batch_norms.append(nn.BatchNorm1d(hidden_dims[i]))
        self.output_layer = nn.Linear(hidden_dims[-1], output_dim)
        nn.init.uniform_(self.output_layer.weight, -2.0, 2.0)
        nn.init.constant_(self.output_layer.bias, 0.0)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, edge_index):
        for i, (conv, bn) in enumerate(zip(self.layers, self.batch_norms)):
            x = conv(x, edge_index)
            x = bn(x)
            x = F.leaky_relu(x, negative_slope=0.2)
            x = self.dropout(x)
        out = self.output_layer(x)
        out = torch.sigmoid(out)
        return out